// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/expression/expr.h>
#include <pollux/expression/string_writer.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/string_encoding_utils.h>
#include <pollux/common/strings/string_impl.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::functions {
    ///  reverse(Array[E]) -> Array[E]
    ///  Takes any array as an input and returns the reversed array.
    ///
    ///  reverse(Varchar) -> Varchar
    ///  Takes any string as an input and returns a string with characters in
    ///  reverse order.
    ///
    ///  reverse(Varbinary) -> Varbinary
    ///  Takes any binary as an input and returns a binary with bytes in reverse
    ///  order.
    class ReverseFunction : public exec::VectorFunction {
    private:
        /// String encoding wrappable function
        template<bool isAscii>
        struct ApplyVarcharInternal {
            static void apply(
                const SelectivityVector &rows,
                const FlatVector<StringView> *input,
                FlatVector<StringView> *result) {
                rows.applyToSelected([&](int row) {
                    auto proxy = exec::StringWriter(result, row);
                    stringImpl::reverse<isAscii>(proxy, input->value_at(row).getString());
                    proxy.finalize();
                });
            }
        };

    public:
        void apply(
            const SelectivityVector &rows,
            std::vector<VectorPtr> &args,
            const TypePtr & /* outputType */,
            exec::EvalCtx &context,
            VectorPtr &result) const override {
            POLLUX_CHECK_EQ(args.size(), 1);

            auto &arg = args[0];

            switch (args[0]->type_kind()) {
                case TypeKind::ARRAY:
                    applyArray(rows, arg, context, result);
                    return;
                case TypeKind::VARCHAR: {
                    const auto ascii = isAscii(arg.get(), rows);
                    applyVarchar(rows, arg, ascii, context, result);
                    return;
                }
                case TypeKind::VARBINARY:
                    // The only difference betwen VARCHAR and VARBINARY input is that
                    // VARBINARY is reversed byte-by-byte, while VARCHAR is reversed
                    // character-by-character. Hence, VARINARY behavior is the same as
                    // VARCHAR with ascii flag set to true.
                    applyVarchar(rows, arg, true /*isAscii*/, context, result);
                    return;
                default:
                    POLLUX_FAIL(
                        "Unsupported input type for 'reverse' function: {}",
                        arg->type()->toString());
            }
        }

        void applyVarchar(
            const SelectivityVector &rows,
            VectorPtr &arg,
            bool isAscii,
            exec::EvalCtx &context,
            VectorPtr &result) const {
            // Capture the pointer to input argument. prepareFlatResultsVector may move
            // it into result.
            auto *originalArg = arg.get();

            prepareFlatResultsVector(result, rows, context, arg, arg->type());
            auto *flatResult = result->as<FlatVector<StringView> >();

            // Input can be constant or flat.
            if (originalArg->is_constant_encoding()) {
                auto value = originalArg->as<ConstantVector<StringView> >()->value_at(0);

                auto proxy = exec::StringWriter(flatResult, rows.begin());
                if (isAscii) {
                    stringImpl::reverse<true>(proxy, value.str());
                } else {
                    stringImpl::reverse<false>(proxy, value.str());
                }
                proxy.finalize();

                auto rawResults = flatResult->mutableRawValues();
                auto reversedValue = rawResults[rows.begin()];

                rows.applyToSelected([&](auto row) { rawResults[row] = reversedValue; });
            } else {
                auto flatInput = originalArg->as<FlatVector<StringView> >();

                StringEncodingTemplateWrapper<ApplyVarcharInternal>::apply(
                    isAscii, rows, flatInput, flatResult);
            }
        }

        bool ensureStringEncodingSetAtAllInputs() const override {
            return true;
        }

        bool propagateStringEncodingFromAllInputs() const override {
            return true;
        }

        void applyArray(
            const SelectivityVector &rows,
            VectorPtr &arg,
            exec::EvalCtx &context,
            VectorPtr &result) const {
            VectorPtr localResult;

            // Input can be constant or flat.
            if (arg->is_constant_encoding()) {
                auto *constantArray = arg->as<ConstantVector<ComplexType> >();
                const auto &flatArray = constantArray->value_vector();
                const auto flatIndex = constantArray->index();

                exec::LocalSingleRow singleRow(context, flatIndex);
                localResult = applyArrayFlat(*singleRow, flatArray, context);
                localResult =
                        BaseVector::wrap_in_constant(rows.end(), flatIndex, localResult);
            } else {
                localResult = applyArrayFlat(rows, arg, context);
            }

            context.moveOrCopyResult(localResult, rows, result);
        }

        VectorPtr applyArrayFlat(
            const SelectivityVector &rows,
            const VectorPtr &vector,
            exec::EvalCtx &context) const {
            auto array_vector = vector->as<ArrayVector>();
            auto elementCount = array_vector->elements()->size();

            // Allocate new vectors for indices.
            auto *pool = context.pool();
            BufferPtr indices = allocate_indices(elementCount, pool);
            auto rawIndices = indices->asMutable<vector_size_t>();

            auto elementsVector = array_vector->elements();
            auto rawSizes = array_vector->rawSizes();
            auto rawOffsets = array_vector->rawOffsets();

            rows.applyToSelected([&](auto row) {
                auto size = rawSizes[row];
                auto offset = rawOffsets[row];

                for (auto i = 0; i < size; ++i) {
                    rawIndices[offset + i] = offset + size - i - 1;
                }
            });

            auto elementsDict =
                    BaseVector::transpose(indices, std::move(elementsVector));

            return std::make_shared<ArrayVector>(
                pool,
                vector->type(),
                array_vector->nulls(),
                rows.end(),
                array_vector->offsets(),
                array_vector->sizes(),
                elementsDict,
                array_vector->get_null_count());
        }

        static std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
            return {
                // array(T) -> array(T)
                exec::FunctionSignatureBuilder()
                .typeVariable("T")
                .returnType("array(T)")
                .argumentType("array(T)")
                .build(),
                // varchar -> varchar
                exec::FunctionSignatureBuilder()
                .returnType("varchar")
                .argumentType("varchar")
                .build(),
                // varbinary -> varbinary
                exec::FunctionSignatureBuilder()
                .returnType("varbinary")
                .argumentType("varbinary")
                .build(),
            };
        }
    };

    POLLUX_DECLARE_VECTOR_FUNCTION(
        udf_reverse,
        ReverseFunction::signatures(),
        std::make_unique<ReverseFunction>());
} // namespace kumo::pollux::functions
